#!/usr/bin/env python3
"""Where the time goes between someone speaking and the call agent acting on it.

  aux-latency [--call <id>] [--transcript <session.jsonl>] [--all]

Defaults to the running call agent's call and session ($AUX_HOME/agent.json). For each line said on the call
it joins three records:

  Tuple       `tuple capture show <call>`: when the utterance ended, and when Tuple finished transcribing it
  feed.log    $AUX_HOME/calls/<call>/feed.log (aux-listen): when the agent's feed command received the line
  session     the agent's Claude Code transcript: when the line was queued for the agent, when its turn
              picked it up, and its first tool call after that

and prints the legs: spoke→tuple (transcription), tuple→feed (Tuple's stream and the pipe), feed→queued (the
Monitor), queued→turn (waiting for the agent's previous turn), turn→act (the agent thinking, for lines it acted on). Lines slower
than 10 s in all are listed one by one (--all lists every line); then p50/p90 per leg.
"""
import argparse
import json
import re
import subprocess
import sys
from datetime import datetime
from pathlib import Path

sys.dont_write_bytecode = True
sys.path.insert(0, str(Path(__file__).resolve().parent))
from aux_context import AUX, cli_path  # noqa: E402

STAMP = re.compile(r"\[(\d{4}-\d\d-\d\dT[^\]]+Z)\]")


def ts(s):
    return datetime.fromisoformat(s.replace("Z", "+00:00")).timestamp()


def tuple_lines(call):
    """{start stamp: (end, finished, speaker, text)} from Tuple's own record of the call."""
    out = {}
    for env in ("staging", "prod", "dev"):
        cli = cli_path(env)
        if not cli:
            continue
        r = subprocess.run([cli, "--format", "json", "capture", "show", call], capture_output=True, text=True, timeout=30)
        for line in r.stdout.splitlines():
            try:
                e = json.loads(line)
            except ValueError:
                continue
            d = e.get("data") or {}
            if e.get("type") == "transcription_finished" and d.get("start"):
                out[d["start"]] = (ts(d["end"]), ts(e["time"]), (d.get("speaker") or "?").split(" ")[0], d.get("text", ""))
        if out:
            break
    return out


def feed_times(call):
    """{start stamp: when the feed received the line}."""
    out = {}
    path = AUX / "calls" / call / "feed.log"
    if path.exists():
        for line in path.read_text(errors="replace").splitlines():
            at, _, text = line.partition("\t")
            m = STAMP.search(text)
            if m:
                out.setdefault(m.group(1), float(at))
    return out


def session_times(path):
    """{start stamp: (queued, picked up, first tool call after)}."""
    queued, picked, acted = {}, {}, {}
    waiting = []
    for line in open(path):
        try:
            d = json.loads(line)
        except ValueError:
            continue
        t = d.get("timestamp")
        if not t:
            continue
        t = ts(t)
        if d.get("type") == "queue-operation" and "Monitor event" in (d.get("content") or ""):
            stamps = STAMP.findall(d["content"])
            if d.get("operation") == "enqueue":
                for s in stamps:
                    queued.setdefault(s, t)
            else:
                for s in stamps:
                    picked.setdefault(s, t)
                waiting += stamps
        elif d.get("type") == "user" and "Monitor event" in line:
            # A notice handed over without a queue entry: picked up when delivered.
            for s in STAMP.findall(line):
                queued.setdefault(s, t)
                picked.setdefault(s, t)
                waiting.append(s)
        elif d.get("type") == "assistant" and waiting:
            blocks = [b for b in d["message"].get("content", []) if isinstance(b, dict)]
            if any(b.get("type") == "tool_use" for b in blocks):
                for s in waiting:
                    acted.setdefault(s, t)
                waiting = []
            elif any(b.get("type") == "text" for b in blocks):
                waiting = []   # a reply with no tool call ("Listening."): those lines needed nothing
    return {s: (queued[s], picked.get(s), acted.get(s)) for s in queued}


def pct(xs):
    xs = sorted(x for x in xs if x is not None)
    if not xs:
        return "     n/a"
    return f"{xs[len(xs) // 2]:5.1f}s / {xs[min(len(xs) - 1, int(len(xs) * 0.9))]:5.1f}s  (n={len(xs)})"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--call")
    ap.add_argument("--transcript")
    ap.add_argument("--all", action="store_true")
    a = ap.parse_args()
    agent = {}
    try:
        agent = json.loads((AUX / "agent.json").read_text())
    except (OSError, ValueError):
        pass
    call = a.call or agent.get("call")
    transcript = a.transcript or agent.get("transcript")
    if not call:
        sys.exit("aux-latency: no call (pass --call)")
    said = tuple_lines(call)
    feed = feed_times(call)
    session = session_times(transcript) if transcript and Path(transcript).exists() else {}
    if not feed:
        print(f"(no feed.log for {call[:8]}: there is no feed log, so tuple→feed and feed→queued are joined)")

    legs = {k: [] for k in ("spoke→tuple", "tuple→feed", "feed→queued", "queued→turn", "turn→act")}
    rows = []
    for start, (end, done, who, text) in sorted(said.items()):
        got = feed.get(start)
        q, p, act = session.get(start, (None, None, None))
        if q is None and got is None:
            continue
        # The catch-up burst when the feed opens replays old lines: not latency.
        if (got or q) - done > 120:
            continue
        leg = {"spoke→tuple": done - end,
               "tuple→feed": (got - done) if got else None,
               "feed→queued": (q - got) if got and q else (q - done if q and not got else None),
               "queued→turn": (p - q) if p and q else None,
               "turn→act": (act - p) if act and p else None}
        for k, v in leg.items():
            if v is not None:
                legs[k].append(v)
        total = (act or p or q or got) - end   # to the action, or to the turn that let it pass
        rows.append((total, start, who, text, leg))

    for total, start, who, text, leg in rows:
        if a.all or total > 10:
            parts = "  ".join(f"{k} {v:4.1f}s" for k, v in leg.items() if v is not None)
            print(f"{start[11:19]} {total:5.1f}s  {who}: {text[:50]!r}\n          {parts}")
    print(f"\n{call[:8]}: {len(rows)} live lines")
    for k, v in legs.items():
        print(f"  {k:12} {pct(v)}")


if __name__ == "__main__":
    main()
